import sets;	
from MatrixIO import readmatrix;
#from mfuncs import *
import mfuncs
import cern;
import jarray
import java.lang.Math as Math
import java.util.Random as rand
from stats import ztest, ttest
import DiskDB
import sets,copy,random
import java.util.Random as Random
from cern.colt.matrix.linalg import DoubleAlgebra as Algebra;
from mfuncs import size
from mfuncs import nonzeros
from mfuncs import solve,solveTranspose,transpose,norm,normp,cor,cov,stddev,abs

import Matrix;
reload(Matrix)
from Matrix import Matrix;

from hep.aida.bin import QuantileDoubleBin1D
from cern.colt.list import DoubleArrayList
import cern.colt.matrix.doublealgo.Sorting as Sorting
import math

import cern.colt.matrix.DoubleFactory2D as DoubleFactory2D

def diag(M,N,v):
	F = DoubleFactory2D.dense
	vec=[]
	q = 1/math.sqrt(N)
	for n in range(0,N):
		vec += [q]
	X = F.diagonal(vec)
	return Matrix(X)

def sampleNet(A):	
	pass

def median(A):
	[M,L]=size(A)
	q = QuantileDoubleBin1D(0.0) 
	x = A.viewColumn(0)
	x = x.toArray()
	q.addAllOfFromTo(DoubleArrayList(x),0,M-1)
	m=q.median()
	q=None
	return m
	
def qnorm(X):
	V=[]
	[M,N] = size(X)
	W = Matrix(M,N)
	for j in range(0,N):
		x = X.viewColumn(j).toArray()
		S = Sorting.mergeSort;
		print "Sorting index: ", j
		v = S.sortIndex(cern.colt.matrix.impl.DenseDoubleMatrix1D(x))
		W[:,j] = X[v,[j]]
		V+=[v]  #keep track of original unsorted indices
	D = diag(M,N,1/N)
	W = W*D
	print "after multiply"
	for j in range(0,N):
		for k in range(0,len(V)):
			v = V[k]
			X[v,j]=W[0:M,j]  # what does this do???
			#for i in range(0,M):  #small v
#				p = W[i,j].get(0,0)
#				X[v[i],j] = p
	print "done"
	return X  ## this is a reference.

def stepwiseNCA(E,Ac,P):
	X = sets.Set()
	[N,M] = size(E)
	not_converged=True
	BIC_new=0
	BIC_old=1
	F = Algebra()
	
	while not_converged:
		for n in range(1,N+1):
			psel=[]
			indexA = [];
			indexA = nonzeros(Ac[n,:]);
			loop_=1
			if len(indexA)>0:
				errList=[]
				while BIC_new < BIC_old:
					for j in indexA:
						indexE=nonzeros(E[n,:]);
						##print "before if"
						if len(indexE)>0:
							##print "in if"
							##print n,j
							##print "a is", Ac[n,j]
							##print "e is", E[n,indexE]
							
							x = F.solveTranspose(E[n,indexE].reshape(1,len(indexE)),P[j,indexE].reshape(1,len(indexE)));
							Ac[n,j] = Matrix(x)

							##print "1a"
							F = Algebra()
							##print Ac[n,j]
							##print P[j,indexE]
							X = Ac*P
							##print E[n,indexE]
							err=E-X;#[n,indexE]-X[n,indexE]
							##print "1c"
							err=norm(err,'fro')
							errList += [err]
					rein=min(errList)
					BIC_old=BIC_new
					BIC_new=-2.0*Math.log(rein**2/M)+Math.log((M+loop_)/(M-loop_-2.0));
					#print BIC_new;
					if (BIC_new < BIC_old):
		   		 	   BIC_old=BIC_new;
		   		 	   loop_=loop_+1;
	   		   		else:
	   		   			  not_converged=False
	   		   			  #print "exiting"
	   		   			  break


	



def convertIDs(E):
	d = DiskDB.Dict('.','probe2gene')
	x = map(lambda x : d[x], E)
	return x

def logRatios(Et,refidx,logForm):
	#print "building log-ratios for E with respect to index: ", refidx 
	E = Matrix(Et)
	x = E.viewColumn(refidx)
	for i in range(0,E.rows()):
		for j in range(0,E.columns()):
			if j==refidx:
				if logForm:
					r = 0
				else:
					r = 1
			else:				
				if logForm:
					r = E.getQuick(i,j)-x.get(i)	
				else:
					r = E.getQuick(i,j)/x.get(i)		
					try:
						r = math.log(r)
					except:
						r = 0
						#print "Error: Cannot take log of zero, ignoring"
			E.setQuick(i,j,r)
	return E

def mergeDataSets(E,Elabels,ExpLabels,E_t,Elabels_t,Explabels_t,progressMan):
	
	# find all Elabels that match Elabels_t
	# check that they are the same size
	idx_t=[]
	for i in Elabels:
		try:
			v = Elabels_t.index(i)
			idx_t += [v]
		except:
			pass
##	#print idx_t
	F = cern.colt.matrix.DoubleFactory2D.dense;
	A = E_t.viewSelection(idx_t,range(0,E_t.columns()))
	B = E
	N = F.appendColumns(B,A)
	E = N
	Elabels=[]
	for i in idx_t:
		Elabels +=[Elabels_t[i]]

	ExpLabels+=Explabels_t

	return (E,Elabels,ExpLabels)
def RankCheck(A):
	(ar,ac)=size(A)
	L=ac
	R=0;
	#print "#...Checking NCA Criteria."
	geneS=sets.Set(range(1,ar))
	tfS=sets.Set(range(1,L+1))
	
	for k in xrange(1,L+1):
		ridx = nonzeros(A[:,k])
		x = sets.Set();
		x.add(k);
		kidx = list(geneS.symmetric_difference(sets.Set(ridx)))
		cidx = list(tfS.symmetric_difference(x))
		R = R + rank(A[kidx,cidx])

	if R < L*(L-1):
		#print "#debug:	Rankcheck failed, NCA cannot be run on this matrix, it is rank deficient, matrix rank is: ", R
		return 0
	else:
		return 1

def Match(E,Elabels,A,Alabels,tfs=None):
	
	# need to startup a progress monitor
	# get total length
	(er,ec)=size(E)
	(ar,ac)=size(A)

	idxA=[];
	idxB=[];
	ok=False
	q=0
	##print len(Elabels),len(Alabels)
	if len(Elabels) < len(Alabels):
		Smaller = Elabels
		Larger = Alabels
	else:
		Smaller = Alabels
		Larger = Elabels

	for x in Smaller:
		try:
			p=Larger.index(x)
			idxB += [p];
			ok=True
			###print "found" 
		except:
			ok=False
			pass

		if ok is True:
			idxA += [q];
			ok=False
		q=q+1

	##print idxA,idxB
	if len(idxA)==0:
		return

	if len(Elabels) < len(Alabels):
		E0=Matrix(len(idxA),ec,0) 
		for i in range(0,len(idxA)):
			for j in range(0,ec):
				E0[i,j] = E[idxA[i],j]
		A0=Matrix(len(idxB),ac,0) 
		for i in range(0,len(idxB)):
			for j in range(0,ac):
				A0[i,j] = A[idxB[i],j]
	else:
		E0=Matrix(len(idxB),ec,0) 
		for i in range(0,len(idxB)):
			for j in range(0,ec):				
				E0[i,j] = E[idxB[i],j]
		A0=Matrix(len(idxA),ac,0) 
		for i in range(0,len(idxA)):
			for j in range(0,ac):
				A0[i,j] = A[idxA[i],j]


	gene_ids = [];
	if len(Elabels) < len(Alabels):
		for x in idxA:
			gene_ids += [Elabels[x]];
	else:
		for x in idxB:
			gene_ids += [Elabels[x]];
			
	A0=Matrix(A0.copy())
	A0.rowNames=copy.deepcopy(gene_ids)
	E0.rowNames=copy.deepcopy(gene_ids)
	A0.columnNames = copy.deepcopy(tfs)
	E0.columnNames = copy.deepcopy(E.columnNames)
	
	
	return (E0,A0,gene_ids,tfs)
	
		
		
def Subnetwork(A,Plabels,tfs):
		#print "#... running Subnetwork TF Selection algorithm."
		(M,L) = size(A)
		#print M,L
		
		rcan=[]
		Ared=abs(A);
		
		temp=[];
		for x in tfs:
			temp += [Plabels[x]]
		Plabels=temp;

		geneS = sets.Set(range(0,M-1))
		tfS = sets.Set(range(0,L-1))
		#print "set is ", tfS
		
		for k in xrange(0,L-1): 
			ridx = nonzeros(Ared[:,k])
			x = sets.Set();
			x.add(k);
			kidx = list(geneS.symmetric_difference(sets.Set(ridx)))
			cidx = list(tfS.symmetric_difference(x))
			s=0
			#print cidx
			for x in cidx:
				B = Ared
				if 1> B.viewSelection(kidx,[x]).zSum():	
					 rcan += [x]
					 #print "coreg"
			
		rcan = sets.Set(rcan);kcan = tfS.symmetric_difference(rcan);
		A0=A[:,list(kcan)];
		tfs=[];
		temp=[];
		for x in kcan:
			temp += [Plabels[x]]

		print "Selected regulators: ",temp
		Plabels=temp;
		return (A0,Plabels)

def gNCA(E,A,P=None,white=None):
		#print "...Running NCA...."
		Ac = Matrix(A.toArray())
		[er,ec]=size(E);
		[ar,ac]=size(A);
		P=Matrix(ac,ec,'r');
#		P.columnNames=E.columnNames
#		P.rowNames=A.rowNames


		epsl_thresh=1e-4
		epsl = 1;
		err=1000;
		err_old=100
		gene_unstable=sets.Set()
		c=0
		##print condE
		while (epsl > epsl_thresh):
			c=c+1
			if c>100:
				c=0
				#print "No convergence restarting..."
				for i in range(0,ar):
					for j in range(0,ac):
						Ac[i,j] = Ac[i,j].get(0) * java.lang.Math.abs(java.lang.Math.random())

			
			for k in xrange(0,ec):
				idxP = nonzeros(P[:,k])
				idxA = nonzeros(Ac[:,idxP])
				if len(idxP)>0:
					P[idxP,k]=solve(Ac[idxA,idxP],E[idxA,k]);

			for n in xrange(0, ar):
			  if n not in gene_unstable:
				indexA = [];
				#print "b4 call"
				indexA = nonzeros(A[n,:]);
				#print "after call"

				if len(indexA)>0:
					indexE=nonzeros(E[n,:]);
					if len(indexE)>0:
						try:
							Ac[n,indexA] = solveTranspose(E[n,indexE],P[indexA,indexE]);
						except:
							Ac[n,indexA] = solveTranspose(P[indexA,indexE],E[n,indexE]);
						 
			err_old=err;
			err=norm(E-Ac*P,'fro')
			epsl=abs(err-err_old)/err;
	

		P._name="TFA Matrix"
		Ac._name="CS Matrix"
#		P.rowNames = A.columnNames
#		P.columnNames = E.columnNames
#		Ac.rowNames = A.rowNames
#		Ac.columnNames = A.columnNames
		#print err
		return (Ac,P,gene_unstable);



def gNCA2(Ec,Ac,P=None,white=None):
		#print ".... Running NCA 2.0 ...."
		A = Ac.copy()
		E = Ec.copy()

		tfs = A.columns()
		exps = E.columns()
		genes = A.rows()

		P = cern.colt.matrix.impl.DenseDoubleMatrix2D(tfs,exps)

		for i in xrange(0,genes):
			for j in range(0,tfs):
				v = A.get(i,j) * Random().nextDouble()
				A.set(i,j,v)

		epsl_thresh=1e-6
		epsl = 1;
		err=1;
		gene_unstable=sets.Set()
		c=0

		F=cern.colt.matrix.linalg.DoubleAlgebra()

		while epsl > epsl_thresh:
			c=c+1
			if c>500:
				c=0
				#print "no convergence, restarting"
				for i in xrange(0,genes):
					for j in range(0,tfs):
						v = A.get(i,j) * Random().nextGaussian()
						A.set(i,j,v)

#				 epsl_thresh=epsl_thresh*10;
			for k in xrange(1,exps):
					V = F.solve(A,E.viewSelection(range(0,genes),[k]))
					r=0
					for tf in range(0,tfs):
							P.set(tf,k,V.get(r,0) )
							r=r+1
					pass
			pass
			for n in xrange(0,genes):
			  if n not in gene_unstable:
				 Am = Matrix(A)
				 Em = Matrix(E)
				 indexA = nonzeros(Am[n+1,:])
				 indexA = map(lambda x: x[1]-1, indexA)
				 indexE = nonzeros(Em[n+1,:])
				 indexE = map(lambda x: x[1]-1,indexE)
				 if len(indexE)>0:
						 Ared = F.solveTranspose(P.viewSelection(indexA,indexE),E.viewSelection([n],indexE))
						 c=0
						 for j in indexA:
							A.set( n,j,Ared.get(c,0) )
							c=c+1
			err_old=err;
			R = Matrix(E)-Matrix(A)*Matrix(P)
			err=norm(R,'fro')
			#print err
			try:
				epsl=abs(err-err_old)/err;
			except:
				epsl=0

			print err


		return (Matrix(A),Matrix(P),gene_unstable);

	
def statsnca(E,A,iterations=None,shuffle=None,Eorig=None):
	
	[genes,tfs] = size(A);
	#print  genes, tfs
	
	iterations=5
	[genes,exps] = size(E);
	#print "dimensions are:", genes, exps

	Eo = Matrix(E)
	
	Ax = cern.colt.matrix.impl.DenseDoubleMatrix3D(iterations,genes,tfs);
	Px = cern.colt.matrix.impl.DenseDoubleMatrix3D(iterations,tfs,exps);
	Er = None
	Ao = Matrix(A.toArray())
	for j in xrange(0,iterations-1):
		if True==shuffle:
			#print "called shuffle"
			[r,c]=size(E)
			[ro,co] = size(Eorig);
			x=range(0,ro)
			random.shuffle(x)
			x=x[0:r]
			###print x
			t = Eorig.viewSelection(x,range(0,co)).copy();
			Er = Matrix(t)

	
		if True==shuffle:
			[Am,Pm,g]=gNCA(Er,A,None);
			[Am,Pm]=normp(Am,Pm)
		else:
			[Am,Pm,g]=gNCA(E,A,None);
			[Am,Pm]=normp(Am,Pm)
		for m in xrange(0,genes-1):
			for n in xrange(0,tfs-1):
				Ax.setQuick(j,m,n,Am.getQuick(m,n));
			for n in xrange(0,tfs-1):
				for l in xrange(0,exps-1):
					Px.setQuick(j,n,l,Pm.getQuick(n,l));


	## now get P_average
	#if shuffle:
	#	return (None,None,None,Ax,Px)
	Pdev = Matrix(tfs,exps);

	for j in xrange(0,tfs-1):
		Z=[]
		idx=[]
		Z = cor( transpose( Matrix(Px.viewRow(j)) ))
		corrV = Z[:,1]
		xvec=[]
	  	for i in range(0,corrV.rows()):
	  		  	###print "inner for"
	  		  	x= corrV[i,0].get(0,0)
	  		  	###print x
			  	if x >= 0.8:
			  		idx += [i]
					xvec += [x]
					x=None
					## need to handle multiply by NEG_1
				
		
		idx = range(0,iterations-1)
		for x in range(0,exps):
			Pj = Matrix(Px.viewSelection(idx,[j],[x]).viewRow(0).copy())
			Pm[j,x] = mfuncs.mean(Pj)
		#	Pdev[j,x] = mfuncs.stddev(Pj)
			
	gene_unstable=sets.Set()
	for n in xrange( 0, genes-1):
			  if n not in gene_unstable:
				indexA = [];
				indexA = nonzeros(A[n,:]);
				if len(indexA)>0:
					indexE=nonzeros(E[n,:]);
					if len(indexE)>0:
						#try:
							Am[n,indexA] = solveTranspose(Pm[indexA,indexE],E[n,indexE]);
						#except:
						 #  gene_unstable.add(n)


		# need to return gene_unstable
	
	[Am,Pm]=normp(Am,Pm)
#	Pm.rowNames=A.columnNames
#	Pm.columnNames=E.columnNames
#	Am.rowNames=A.rowNames
#	Am.columnNames=A.columnNames
	Pm.varname='p'
	
	R = E-Am*Pm
	R = norm(R,'fro')
	#print R
	if shuffle:
		print "Warning: During Shuffling Procedure, the following genes have unstable CS values: ", gene_unstable
	return (Am,Pm,Pdev,Ax,Px)
	

def FastNCA(E,A):
	pass

def setKnockOutMatrix(P):
	pass

def sigTFs(q,prx):
	
	if isinstance(q,Matrix):
		px = q
	else:
		px = q 
	r = px.rows()
	c = px.columns()
	sigTF = Matrix(r,c,0)
	for i in xrange(0,r):
		for j in xrange(0,c):
			p = px.viewRow(i).get(j)
			q = prx.viewRow(i).viewColumn(j)
			h = ttest(p,q,0.05)
			sigTF.set(i,j,h)
	return sigTF
